import abc
from typing import Generic, TypeVar

from .task_stepping import SteppingTask, SteppingTaskPlugin

AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')


class ProcessingTaskPlugin(SteppingTaskPlugin, Generic[AnyInput, AnyOutput], abc.ABC):
    def on_before_process(self, input: AnyInput) -> AnyInput:
        return input

    def on_after_process(self, output: AnyOutput) -> AnyOutput:
        return output


class ProcessingTask(SteppingTask, Generic[AnyInput, AnyOutput], abc.ABC):
    def _step(self):
        input = self._produce()
        input = self._before_process(input)
        output = self._process(input)
        output = self._after_process(output)
        self._consume(output)

    @abc.abstractmethod
    def _produce(self) -> AnyInput:
        pass

    def _before_process(self, input: AnyInput) -> AnyInput:
        return self.dock.scan_agg_plugins(
            self.dock.iter_plugins(ProcessingTaskPlugin), input,
            lambda plugin, input: plugin.on_before_process(input))

    @abc.abstractmethod
    def _process(self, input: AnyInput) -> AnyOutput:
        pass

    def _after_process(self, output: AnyOutput) -> AnyOutput:
        return self.dock.scan_agg_plugins(
            self.dock.iter_plugins(ProcessingTaskPlugin), output,
            lambda plugin, output: plugin.on_after_process(output))

    @abc.abstractmethod
    def _consume(self, output: AnyOutput):
        pass
